# -*- coding: utf-8 -*-
"""
@author:xuyuntao
@time:2021/3/4:9:10
@email:xuyuntao@189.cn
"""
import numpy
import matplotlib.font_manager
import time
import os
from Generator import *
from tqdm import tqdm
from charactersAlgorithm import *
from Tools import MT,get_desktop
import Tools.drawing as Draw
import threading
import argparse
parser = argparse.ArgumentParser(description='控制性参数')
parser.add_argument('--cfgNum', type=int, default=0)
parser.add_argument('--ifSave', type=str, default="True")
args = parser.parse_args()

"""用于绘制第三章、第四章中所有特征值随信噪比变化图，
每幅图中绘制的信号种类、特征值、信噪比轴范围均通过ValueShowCfg文件夹中的配置文件定义
注意：该程序为需要在log窗口输入运行程序号的版本
可以通过runfile或os.system外部传入控制参数：
cfgNum：绘图配置号，若为0则需要手工console窗口输入参数，否则使用给定参数
"""
# ----------------------------参数初始化----------------------------
# 以防下面cfg文件中没有特定变量
lineWidth=1   # 绘图线宽
figSizeNum = 8  # 窗口大小基本量
figSize = (figSizeNum, round(figSizeNum / 4 * 3))  # 窗口大小
figScale = (15, round(15 / 4 * 3))  # 窗口大小
fontSize = 21  # 其他字体大小
labelSize = 18  # 轴刻度字号
legendSize=fontSize  # 图例字号
numbersInsideSize=8  # 内部标注数字字号
drawTypes=[0,1]   # 0表示热力图，1表示折线图
marker=True   # 折线图显示marker，为list可指定各折线显示的marker，True则使用默认显示序列
# ---调制类型轴---
modulationTypes=[MT("ask",2),MT("ask",4),\
                 MT("psk",2),MT("psk",4),\
                 MT("qam",4),MT("qam",16),\
                 MT("fsk",2),MT("fsk",4)]    # 按这个循环计算各调制种类的高阶累积量
SNRList=[_ for _ in range(10,-5,-1)]  # 信噪比轴  [-1,0,1,2,3,...]
varList=[R_avg, Gamma_max_avg, Sigma_aa_avg, Sigma_ap_avg, Sigma_af_avg] # 计算特征值列表
# 计算几种不同高阶累积量  [func_1,func_2,...]
specialList=[Sigma_ap_avg, Sigma_af_avg]  # 需要传入其他参数的特征值函数
fskSpecial=fskC41C20  # fsk变换函数需要特别考虑，传入sep参数
sampleLength=[300]   # 取几种不同采样长度
norm=True   # 是否进行归一化
SignalOrConstel=0   # 0为调制信号Signal，1为基带信号Constel
awgnOrAntArr=0   # 是否使用阵列天线，0为使用普通天线，1为使用阵列天线
signalWay=30   # 使用阵列天线时设定的信号来源角度
# ---其他参数---
a_t=1   # 论文at
# ---码元数据---
avgTimes=50   # 一次生成的信号数，用于计算平均特征值
bytesRate=16  # 码元速率
carrierFreq=32  # 载波速率
sampleFreq=60*carrierFreq  # 采样速率
seperate = 16  # fsk两频率相较载波的偏差
fsk_phase_continue = True  # 是否相位连续

# ----------------------------控制参数传入----------------------------
argsIn=False
if args.cfgNum==0:
    ifSaveFiles = bool(input("是否保存？(y/n)").upper() == "Y")  # False
    configurationNum=int(input("输入配置程序序号："))
else:
    argsIn=True
    ifSaveFiles=eval(args.ifSave)
    configurationNum=args.cfgNum
exec("from ValueShowCfg.cfg_{} import *".format(configurationNum))

# ----------------------------数据计算部分----------------------------
# ---绘图控制---
startT=time.time()
workPath=os.path.join(get_desktop(),"Paper_Pictures\\","varResults_avg\\")  # 保存文件路径，末尾必须有\\，且全部只出现\\
if os.path.isdir(workPath):
    pass
else:
    os.mkdir(workPath)
saveDPI=400    # 除星座图的保存文件dpi
fileType="svg"   # 除星座图外的保存文件格式
ifSave=ifSaveFiles   # 是否保存图像
# dBAxis=False  # 是否以dB为轴单位
colormap="nipy_spectral"  # colormap
zhFont=matplotlib.font_manager.FontProperties(fname='C:\Windows\Fonts\simsun.ttc')  # 中文字体
# plt.rcParams['xtick.direction'] = 'in'
# plt.rcParams['ytick.direction'] = 'in'
# ---信号生成---
# startTime=time.time()
SignalList=[]  # 信号缓冲区，第一层比特长度，第二层调制类型，第三层平均次数，第四层几种信号
with tqdm(total=len(sampleLength)*avgTimes*len(modulationTypes), desc="生成调制信号") as pbar:
    for sl in sampleLength:  # 创建多个类型的数字调制产生器，后面循环调用
        modSignals=[]
        for mt in modulationTypes:
            avgSignals=[]
            for at in range(avgTimes):
                bitsData = getBitsData(mt.step, sl)
                DMG = DigitalModulationGenerator(bitsData,
                                                 bytesRate, carrierFreq, sampleFreq)
                DMG.setFSKSeperateFreq(seperate)
                cst=DMG.get_Constel_Signal(mt.type,mt.step,True)
                avgSignals.append(cst)
                pbar.update(1)
            modSignals.append(avgSignals)
        SignalList.append(modSignals)
        del modSignals,avgSignals,cst
# endTime = time.time()
# print("生成信号共耗时：", str(datetime.timedelta(seconds=abs(endTime - startTime))))
del sl,mt,at#,endTime,startTime

def calculate(awgnOrAntArr,SignalOrConstel,
              signals_arr,constels_arr,timeSeries_arr,
              signalWay_arr,snr,snrN,varList,specialList,
              sl,mT, carrierFreq, a_t):
    global vars
    if awgnOrAntArr:  # antArr
        signals_noise = antArrChan(signals_arr, snr, signalWay_arr)
        constels_noise = antArrChan(constels_arr, snr, signalWay_arr)
    else:
        signals_noise = awgn(signals_arr, snr)
        constels_noise = awgn(constels_arr, snr)
    if SignalOrConstel:  # constel
        del signals_noise
        for varN in range(len(varList)):
            varFunc = varList[varN]
            lock.acquire()
            if varFunc in specialList:
                vars[sl, mT, snrN, varN] = varFunc(constels_noise, timeSeries_arr, carrierFreq, a_t, 0.1 * numpy.pi)
            else:
                if varFunc==fskSpecial:
                    vars[sl, mT, snrN, varN] = varFunc(constels_noise,timeSeries_arr,seperate)
                else:
                    vars[sl, mT, snrN, varN] = varFunc(constels_noise)
            lock.release()
    else:  # signal
        del constels_noise
        for varN in range(len(varList)):
            varFunc = varList[varN]
            lock.acquire()
            if varFunc in specialList:
                vars[sl, mT, snrN, varN] = varFunc(signals_noise, timeSeries_arr, carrierFreq, a_t, 0.1 * numpy.pi)
            else:
                if varFunc==fskSpecial:
                    vars[sl, mT, snrN, varN] = varFunc(signals_noise,timeSeries_arr,seperate)
                else:
                    vars[sl, mT, snrN, varN] = varFunc(signals_noise)
            lock.release()

vars=numpy.zeros([len(sampleLength),len(modulationTypes),len(SNRList),len(varList)], dtype=numpy.complex128)
lock=threading.RLock()
# 高阶累积量缓冲区
# 信号缓冲区，第一层比特长度，第二层调制类型，第三层平均次数，第四层几种信号
with tqdm(total=len(sampleLength)*len(modulationTypes)*len(SNRList), desc="各变量计算") as pbar:
    for sl in range(len(sampleLength)):  # 按抽样点数进行第一级循环

        # signals=[]   # [(signal_1,time_1),(signal_2,time_2),...]
        # C21NormList = []
        for mT in range(len(modulationTypes)):
            signals=[]
            constels=[]
            timeSeriesList=[]
            for at in range(avgTimes):
                signals.append(SignalList[sl][mT][at][1])
                constels.append(SignalList[sl][mT][at][0])
                timeSeriesList.append(SignalList[sl][mT][at][2])
            signals_arr=numpy.array(signals)
            constels_arr=numpy.array(constels)
            timeSeries_arr=numpy.array(timeSeriesList)
            signalWay_arr = numpy.ones(signals_arr.shape[0], dtype=numpy.float32) * signalWay
            del signals,constels,timeSeriesList,at
            threadList=[]
            for snrN in range(len(SNRList)):
                snr=SNRList[snrN]

                threadList.append(threading.Thread(target=calculate,args=[awgnOrAntArr, SignalOrConstel,
                                                   signals_arr, constels_arr, timeSeries_arr,
                                                   signalWay_arr, snr, snrN, varList, specialList,
                                                   sl, mT, carrierFreq, a_t]))
                threadList[-1].setDaemon(True)
                threadList[-1].start()
            for th in threadList:
                th.join()
                pbar.update(1)
            del signals_arr,constels_arr,timeSeries_arr,signalWay_arr,snrN


# # -------------------绘图-------------------
MTLabel=[str(_[1])+_[0].upper() for _ in modulationTypes]  # 调制类型转换为标签
figures=[]  # mpl.fig缓冲区
figNum=1 # fig编号
SNR_Arr=numpy.array(SNRList).reshape([1,-1]).repeat(len(modulationTypes),axis=0)
with tqdm(total=len(sampleLength)*len(varList)*len(drawTypes),desc="绘图") as pbar:
    for sl in range(len(sampleLength)):  # 按抽样点数进行第一级循环
        for varN in range(len(varList)):  # 特征值函数
            data=numpy.abs(vars[sl,:,:,varN])
            if 1 in drawTypes:  # 折线图
                fig=Draw.drawLine(figNum,SNR_Arr,data,None,None,MTLabel,marker=marker,
                                  figName="[cfg{4}]_[Fig{0}]_[{1}b]_[{2}]_[{3}]_[line]".
                                  format(figNum,sampleLength[sl],
                                         "antArr" if awgnOrAntArr else "awgn",
                                         varList[varN].__name__,configurationNum),
                                  figScale=figSize)
                figNum+=1
                figures.append(fig)
                pbar.update(1)
            if 0 in drawTypes:
                fig=Draw.drawColorMap(figNum,data,None,None,SNRList,MTLabel,
                                     figName="[cfg{4}]_[Fig{0}]_[{1}b]_[{2}]_[{3}]_[colormap]".
                                     format(figNum, sampleLength[sl],
                                            "antArr" if awgnOrAntArr else "awgn",
                                            varList[varN].__name__,configurationNum),
                                      figScale=figScale)
                figNum+=1
                figures.append(fig)
                pbar.update(1)

import matplotlib.pyplot as plt
with tqdm(total=len(figures),desc="显示+保存") as pbar:
    for fig in figures:  # 保存+显示
        figName=fig.canvas.manager.get_window_title()
        fig.show()
        if ifSave:
            fig.savefig(workPath+figName+".{0}".format(fileType),dpi=saveDPI,format=fileType,bbox_inches='tight',pad_inches=0.0,transparent=True)
        if argsIn:
            fig.clf()
            fig.clear()
            plt.close(fig)
        pbar.update(1)
